import { RoleMessageType, RoleType } from "../enum";
import { ChatMessage } from "../message/chatMessage";
import { SystemMessage } from "../message/systemMessage";
import { getPromptFromKey, promptTemplateDict } from "../prompt";
import { ChatAgent, RawChatAgentConfig, ChatAgentResponse } from "./chatAgent";

/**
 * 将任务指定为更具体的任务的Agent
 */
export class TaskSpecifyAgent extends ChatAgent {
    constructor(config: RawChatAgentConfig) {
        const systemMessage = new SystemMessage({
            roleName: "TaskSpecify",
            roleType: RoleType.Assistant,
            roleMessageType: RoleMessageType.AgentSystem,
            content: "你可以将任务指定为更具体的任务。"
        });
        super({ ...config, systemMessage });
    }

    async step(inputMessage: ChatMessage): Promise<ChatAgentResponse> {
        this.reset();
        const prompt = getPromptFromKey(
            promptTemplateDict.baseTemplate.taskSpecify,
            this.language
        ).format({
            assistantRole: this.roleName,
            userRole: inputMessage.roleName,
            task: inputMessage.content,
            wordLimit: String(this.modelTokenLimit)
        });
        const taskMessage = new ChatMessage({
            roleName: this.roleName,
            roleType: this.roleType,
            roleMessageType: RoleMessageType.AgentReceiving,
            content: prompt
        });
        return await super.step(taskMessage);
    }
}

/**
 * 将任务分解为子任务的Agent
 */
export class TaskPlannerAgent extends ChatAgent {
    constructor(config: RawChatAgentConfig) {
        const systemMessage = new SystemMessage({
            roleName: "TaskPlanner",
            roleType: RoleType.Assistant,
            roleMessageType: RoleMessageType.AgentSystem,
            content: "你可以将任务分解为子任务。"
        });
        super({ ...config, systemMessage });
    }

    async step(inputMessage: ChatMessage): Promise<ChatAgentResponse> {
        this.reset();
        const prompt = getPromptFromKey(
            promptTemplateDict.baseTemplate.taskPlanner,
            this.language
        ).format({
            task: inputMessage.content
        });
        const taskMessage = new ChatMessage({
            roleName: this.roleName,
            roleType: this.roleType,
            roleMessageType: RoleMessageType.AgentReceiving,
            content: prompt
        });
        return await super.step(taskMessage);
    }
}
